{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ec7f2e31",
   "metadata": {
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/clutter.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af186df",
   "metadata": {
    "id": "A4QOaw_zYLfI",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import Rgba, StartMeshcat\n",
    "from pydrake.multibody.benchmarks import MassDamperSpringAnalyticalSolution\n",
    "\n",
    "from manipulation import running_as_notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03cfcd00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97343417",
   "metadata": {
    "id": "7aGJ9e3lJlF8"
   },
   "source": [
    "# Let's explore stiffness\n",
    "\n",
    "As \"spring-mass-damper\" system is enough to make the basic point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a67a3a",
   "metadata": {
    "id": "7dimkJhd9DjP"
   },
   "outputs": [],
   "source": [
    "def spring_mass_damper_example():\n",
    "    meshcat.Delete()\n",
    "    meshcat.DeleteAddedControls()\n",
    "    meshcat.SetProperty(\"/Background\", \"visible\", False)\n",
    "    meshcat.Set2dRenderMode(xmin=-2, xmax=2, ymin=-2, ymax=2)\n",
    "\n",
    "    meshcat.AddSlider(name=\"mass\", min=0.02, max=10, value=1.0, step=0.1)\n",
    "    meshcat.AddSlider(name=\"b\", min=0.02, max=10, value=1.0, step=0.1)\n",
    "    meshcat.AddSlider(\n",
    "        name=\"k\",\n",
    "        min=0.02,\n",
    "        max=10,\n",
    "        value=1.0,\n",
    "        step=0.1,\n",
    "        decrement_keycode=\"ArrowLeft\",\n",
    "        increment_keycode=\"ArrowRight\",\n",
    "    )\n",
    "    meshcat.AddSlider(\n",
    "        name=\"h\",\n",
    "        min=0.02,\n",
    "        max=10,\n",
    "        value=0.1,\n",
    "        step=0.1,\n",
    "        decrement_keycode=\"ArrowDown\",\n",
    "        increment_keycode=\"ArrowUp\",\n",
    "    )\n",
    "\n",
    "    meshcat.AddButton(\"Stop Demo\", \"Escape\")\n",
    "\n",
    "    N = 200\n",
    "    tf = 20\n",
    "    ts = np.linspace(0, tf, N)\n",
    "    xs = np.zeros((N, 2))\n",
    "    Xs = np.zeros((3, N))\n",
    "\n",
    "    Nd = 200\n",
    "    xd = np.zeros((Nd, 2))\n",
    "    Xd = np.zeros((3, Nd))\n",
    "\n",
    "    prev = None\n",
    "\n",
    "    while meshcat.GetButtonClicks(\"Stop Demo\") < 1:\n",
    "        mass = meshcat.GetSliderValue(\"mass\")\n",
    "        b = meshcat.GetSliderValue(\"b\")\n",
    "        k = meshcat.GetSliderValue(\"k\")\n",
    "        h = meshcat.GetSliderValue(\"h\")\n",
    "        if [mass, b, k, h] == prev:\n",
    "            continue\n",
    "        prev = [mass, b, k, h]\n",
    "\n",
    "        sol = MassDamperSpringAnalyticalSolution(mass, b, k)\n",
    "        sol.SetInitialValue(1, 0)\n",
    "\n",
    "        for i in range(N):\n",
    "            xs[i] = [sol.get_x(ts[i]), sol.get_xDt(ts[i])]\n",
    "        Xs[0, :] = xs[:, 0]\n",
    "        Xs[2, :] = xs[:, 1]\n",
    "        meshcat.SetLine(\"analytical\", Xs, line_width=4, rgba=Rgba(0, 0, 1, 1))\n",
    "\n",
    "        xd[0] = [1, 0]\n",
    "        Ad = np.eye(2) + h * np.array([[0, h], [-k * h / mass, -b * h / mass]])\n",
    "        for i in range(1, Nd):\n",
    "            xd[i] = Ad @ xd[i - 1]\n",
    "        Xd[0, :] = xd[:, 0]\n",
    "        Xd[2, :] = xd[:, 1]\n",
    "        meshcat.SetLine(\"euler\", Xd, line_width=4, rgba=Rgba(1, 0, 0, 1))\n",
    "\n",
    "        if not running_as_notebook:\n",
    "            break\n",
    "\n",
    "    meshcat.DeleteButton(\"Stop Demo\")\n",
    "\n",
    "\n",
    "spring_mass_damper_example()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
